"""
涨跌幅计算器模块

功能：基于上一交易日收盘价计算当前涨跌幅，实现交易控制逻辑

使用方法：
```python
from xtquant import xtdata
from 涨跌幅 import UpDownCalculator

# 初始化
calculator = UpDownCalculator(xtdata)

# 获取涨跌幅信息（传入当前价格）
current_price = 10.25
result = calculator.get_updown("520500.SH", current_price)

# 检查是否可以交易
if result['can_trade']:
    print(f"可以交易，建议RSI阈值: {result['adjusted_rsi']}")
else:
    print(f"暂停交易，原因: {result['reason']}")
```

返回结果：
- updown: 当前涨跌幅百分比
- can_trade: 是否可以交易
- adjusted_rsi: 调整后的RSI阈值
- reason: 交易控制原因
- prev_close_price: 上一交易日收盘价
- current_price: 当前价格
"""

from xtquant import xtdata
import pandas as pd
import time
from datetime import datetime
from typing import Dict, Optional

class UpDownCalculator:
    """
    涨跌幅计算器
    实现功能：基于上一交易日收盘价计算涨跌幅，当前跌破1%时stock_min_rsi设置为5，跌破2%停止交易，回到1.5%持续10分钟后继续交易
    """
    
    def __init__(self, xtdata=None):
        self.xtdata = xtdata
        # 存储每只股票的上一交易日收盘价格
        self.prev_close_prices = {}
        # 存储每只股票的跌幅状态
        self.drop_status = {}
        # 交易时间段常量
        self.MORNING_START = datetime.strptime("09:30", "%H:%M").time()
        self.MORNING_END = datetime.strptime("11:30", "%H:%M").time()
        self.AFTERNOON_START = datetime.strptime("13:00", "%H:%M").time()
        self.AFTERNOON_END = datetime.strptime("15:00", "%H:%M").time()
    
    def _get_prev_close_price(self, stock_code: str) -> Optional[float]:
        """获取股票的上一交易日收盘价格"""
        try:
            # 如果已经获取过，直接返回
            if stock_code in self.prev_close_prices:
                return self.prev_close_prices[stock_code]
            
            # 获取最近5个交易日的数据，确保能获取到上一交易日
            end_date = datetime.now().strftime("%Y%m%d")
            start_date = (datetime.now() - pd.Timedelta(days=5)).strftime("%Y%m%d")
            
            # 下载日线数据
            self.xtdata.download_history_data(
                stock_code=stock_code,
                period="1d",
                start_time=start_date,
                end_time=end_date
            )
            time.sleep(0.1)
            
            # 获取数据
            data = self.xtdata.get_local_data(
                stock_list=[stock_code],
                period='1d',
                start_time=start_date,
                end_time=end_date
            )
            
            df = data.get(stock_code) if data else None
            if df is None or df.empty:
                return None
            
            # 识别时间列和收盘价列
            time_col = next((c for c in df.columns if 'time' in c.lower()), None)
            close_cols = ['close', 'lastPrice', 'price']
            close_col = next((c for c in close_cols if c in df.columns), None)
            
            if not time_col or not close_col:
                return None
            
            # 时间格式统一处理
            if str(df[time_col].dtype) in ('int64', 'float64'):
                df['__time__'] = pd.to_datetime(df[time_col], unit='ms') + pd.Timedelta(hours=8)
            else:
                try:
                    df['__time__'] = pd.to_datetime(df[time_col], format='%Y%m%d%H%M%S', errors='coerce')
                except Exception:
                    df['__time__'] = pd.to_datetime(df[time_col], errors='coerce')
            
            df = df[df['__time__'].notna()].sort_values('__time__')
            
            # 获取上一交易日的收盘价（倒数第二个交易日，因为最后一个可能是今天）
            if len(df) >= 2:
                # 获取倒数第二个交易日的收盘价
                prev_close_price = df.iloc[-2].get(close_col)
                if prev_close_price and prev_close_price > 0:
                    self.prev_close_prices[stock_code] = float(prev_close_price)
                    return float(prev_close_price)
            elif len(df) == 1:
                # 如果只有一天数据，使用当天的收盘价
                prev_close_price = df.iloc[-1].get(close_col)
                if prev_close_price and prev_close_price > 0:
                    self.prev_close_prices[stock_code] = float(prev_close_price)
                    return float(prev_close_price)
            
            return None
            
        except Exception as e:
            print(f"⚠️ 获取{stock_code}上一交易日收盘价格失败: {e}")
            return None
    
    def get_updown(self, stock_code: str, current_price: float = None) -> Dict:
        """
        获取涨跌幅信息并检查交易控制状态
        返回包含交易控制信息的字典
        """
        try:
            # 获取上一交易日收盘价格
            prev_close_price = self._get_prev_close_price(stock_code)
            if prev_close_price is None:
                return {
                    "updown": 0.0,
                    "can_trade": True,
                    "adjusted_rsi": 10,
                    "reason": "无法获取上一交易日收盘价格"
                }
            
            # 如果没有提供当前价格，尝试从xtdata获取实时价格
            if current_price is None:
                try:
                    # 方法1：尝试获取实时tick数据
                    tick_data = self.xtdata.get_market_data([stock_code], period='tick', count=1)
                    if tick_data and stock_code in tick_data and tick_data[stock_code]:
                        # 获取最新的tick数据
                        latest_tick = tick_data[stock_code][0] if isinstance(tick_data[stock_code], list) else tick_data[stock_code]
                        current_price = latest_tick.get('lastPrice')
                        if current_price and current_price > 0:
                            print(f"📊 从tick数据获取{stock_code}当前价格: {current_price}")
                        else:
                            current_price = None
                    
                    # 方法2：如果tick数据获取失败，尝试获取分钟数据
                    if current_price is None:
                        minute_data = self.xtdata.get_market_data([stock_code], period='1m', count=1)
                        if minute_data and stock_code in minute_data and minute_data[stock_code]:
                            latest_minute = minute_data[stock_code][0] if isinstance(minute_data[stock_code], list) else minute_data[stock_code]
                            current_price = latest_minute.get('close') or latest_minute.get('lastPrice')
                            if current_price and current_price > 0:
                                print(f"📊 从分钟数据获取{stock_code}当前价格: {current_price}")
                    
                    # 如果都获取失败，使用上一交易日收盘价
                    if current_price is None or current_price <= 0:
                        current_price = prev_close_price
                        print(f"⚠️ 无法获取{stock_code}实时价格，使用上一交易日收盘价: {prev_close_price}")
                        
                except Exception as e:
                    print(f"⚠️ 获取{stock_code}实时价格失败: {e}")
                    current_price = prev_close_price
            
            # 计算当前涨跌幅（相对于上一交易日收盘价）
            current_drop_pct = ((current_price - prev_close_price) / prev_close_price) * 100
            current_time = time.time()
            
            # 初始化股票状态
            if stock_code not in self.drop_status:
                self.drop_status[stock_code] = {
                    "drop_1_percent_time": None,
                    "drop_2_percent_time": None,
                    "recovery_1_5_percent_time": None,
                    "trading_suspended": False,
                    "recovery_duration": 0
                }
            
            status = self.drop_status[stock_code]
            
            # 检查是否跌破1%
            if current_drop_pct <= -0 and status["drop_1_percent_time"] is None:
                status["drop_1_percent_time"] = current_time
                print(f"⚠️ {stock_code} 跌破1%，跌幅: {current_drop_pct:.2f}%")
            
            # 检查是否跌破2%
            if current_drop_pct <= -0.2 and status["drop_2_percent_time"] is None:
                status["drop_2_percent_time"] = current_time
                status["trading_suspended"] = True
                print(f"🚫 {stock_code} 跌破2%，暂停交易，跌幅: {current_drop_pct:.2f}%")
            
            # 检查是否回到1.5%
            if current_drop_pct >= -1.5:
                if status["recovery_1_5_percent_time"] is None:
                    status["recovery_1_5_percent_time"] = current_time
                    status["recovery_duration"] = 0
                    print(f"📈 {stock_code} 回到1.5%，跌幅: {current_drop_pct:.2f}%")
            else:
                # 如果又跌破1.5%，重置恢复时间
                status["recovery_1_5_percent_time"] = None
                status["recovery_duration"] = 0
            
            # 计算回到1.5%后的持续时间
            if status["recovery_1_5_percent_time"] is not None:
                recovery_duration = (current_time - status["recovery_1_5_percent_time"]) / 60  # 转换为分钟
                status["recovery_duration"] = recovery_duration
                
                # 如果持续10分钟以上，恢复交易
                if recovery_duration >= 10 and status["trading_suspended"]:
                    status["trading_suspended"] = False
                    print(f"✅ {stock_code} 回到1.5%持续{recovery_duration:.1f}分钟，恢复交易")
            
            # 根据跌幅动态调整RSI阈值
            adjusted_rsi = 10  # 默认RSI阈值
            if current_drop_pct <= -1.5:
                adjusted_rsi = 5
                print(f"📉 {stock_code} 跌破1%，RSI阈值调整为5，当前跌幅: {current_drop_pct:.2f}%")
            
            # 返回涨跌幅和交易控制信息
            return {
                "updown": current_drop_pct,
                "can_trade": not status["trading_suspended"],
                "adjusted_rsi": adjusted_rsi,
                "reason": f"跌幅过大({current_drop_pct:.2f}%)，暂停交易" if status["trading_suspended"] else "跌幅在允许范围内",
                "recovery_duration": status["recovery_duration"],
                "prev_close_price": prev_close_price,
                "current_price": current_price
            }
            
        except Exception as e:
            print(f"⚠️ 计算{stock_code}涨跌幅失败: {e}")
            return {
                "updown": 0.0,
                "can_trade": True,
                "adjusted_rsi": 10,
                "reason": "计算失败，允许交易"
            }
    
    def reset_status(self, stock_code: str = None):
        """重置跌幅状态"""
        if stock_code:
            if stock_code in self.drop_status:
                self.drop_status[stock_code] = {
                    "drop_1_percent_time": None,
                    "drop_2_percent_time": None,
                    "recovery_1_5_percent_time": None,
                    "trading_suspended": False,
                    "recovery_duration": 0
                }
            if stock_code in self.prev_close_prices:
                del self.prev_close_prices[stock_code]
        else:
            self.drop_status.clear()
            self.prev_close_prices.clear()
        print(f"🔄 {stock_code or '所有股票'} 价格跌幅状态已重置")
    
    def get_status(self, stock_code: str = None):
        """获取当前状态信息"""
        if stock_code:
            if stock_code in self.drop_status:
                status = self.drop_status[stock_code].copy()
                status["stock_code"] = stock_code
                status["prev_close_price"] = self.prev_close_prices.get(stock_code)
                return status
            return None
        else:
            return {code: self.get_status(code) for code in self.drop_status.keys()}


# 使用示例
if __name__ == "__main__":
    from xtquant import xtdata
    
    # 初始化涨跌幅计算器
    calculator = UpDownCalculator(xtdata)
    
    # 示例股票代码
    stock_code = "159605.SZ"
    
    print("=== 简单使用示例 ===")
    print(f"股票代码: {stock_code}")
    print()
    
    # 示例：传入当前价格
    current_price = 1.271  # 假设当前价格
    result = calculator.get_updown(stock_code, current_price)
    
    print("涨跌幅信息:")
    print(f"  当前价格: {current_price}")
    print(f"  当前涨跌幅: {result['updown']:.2f}%")
    print(f"  是否可交易: {'是' if result['can_trade'] else '否'}")
    print(f"  建议RSI阈值: {result['adjusted_rsi']}")
    print(f"  上一交易日收盘价: {result.get('prev_close_price', 'N/A')}")
    
    if not result['can_trade']:
        print(f"  暂停原因: {result['reason']}")
    
    print()
    print("=== 在行情监控中使用 ===")
    print("在tick回调中直接调用:")
    print("current_price = data0.get('lastPrice')  # 从tick数据获取当前价格")
    print("updown_info = calculator.get_updown(stock_code, current_price)")
    print("if not updown_info['can_trade']:")
    print("    return  # 暂停交易")
    print("rsi_threshold = updown_info['adjusted_rsi']  # 使用调整后的RSI阈值")